package org.example.doris;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.DruidPooledConnection;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.time.StopWatch;
import org.apache.http.HttpEntity;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPut;
import org.apache.http.entity.ByteArrayEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.DefaultRedirectStrategy;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

@Slf4j
public class TagHttpTask implements Runnable {

    private DruidDataSource dataSource;

    private static final AtomicLong userStart = new AtomicLong(320281199508233775l);

    private static final Long userSize = 6000000l;

    private static final Integer tagSize = 250;

    private static final List<String> userList = new ArrayList<>();

    private static final List<String> tagList = new ArrayList<>();

    private static final Integer batchSize = 10000;

    private static CountDownLatch latch = null;

    private static final StopWatch stopwatch = new StopWatch();

    private static final String loadUrl = "http://192.168.186.66:8040/api/dorisdb_test/ep_taglist/_stream_load";

    private static final String username = "test_user";

    private static final String password = "11111";

    public TagHttpTask(DruidDataSource dataSource, CountDownLatch latch) {
        this.dataSource = dataSource;
        this.latch = latch;

        // 用户身份证
        stopwatch.start();
        for (int i = 0; i < userSize; i++) {
            userList.add(String.valueOf(userStart.getAndIncrement()));
        }
        stopwatch.stop();
        log.info("初始化用户耗时：{}毫秒", stopwatch.getTime());

        // 标签名字
        stopwatch.reset();
        stopwatch.start();
        for (int i = 1; i <= tagSize; i++) {
            tagList.add("tag" + i);
        }
        stopwatch.stop();
        log.info("初始化标签耗时：{}毫秒", stopwatch.getTime());
    }

    @Override
    public void run() {

        // 每次10万条记录
        long batchNum = userSize / batchSize;
        try {

            // 250个标签
            for (int i = 0; i < tagSize; i++) {

                for (int j = 0; j < batchNum; j++) {
                    // 每次1万条提交一次
                    stopwatch.reset();
                    stopwatch.start();
                    List<byte[]> list = new ArrayList<>();
                    int totalBytes = 0;
                    for (int k = 0; k < batchSize; k++) {
                        String content = String.format("%d,%s,%s", i, tagList.get(i), userList.get(k + j * batchSize));
                        byte[] conBytes = content.getBytes();
                        list.add(conBytes);
                        totalBytes = totalBytes + conBytes.length;
                    }

                    // 发送数据
                    sendData(list, totalBytes);
                    stopwatch.stop();
                    log.info("{}条耗时：{}", batchSize, stopwatch.getTime());
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
//        List<byte[]> list = new ArrayList<>();
//        byte[] record = "4,tag2,320281199508233775".getBytes();
//        list.add(record);
//        sendData(list, record.length);
        latch.countDown();
    }

    public void sendData(List<byte[]> list, int totalBytes) {
        byte[] lineBytes = "\r\n".getBytes();
        ByteBuffer buffer = ByteBuffer.allocate(totalBytes + (list.size()) * lineBytes.length);

        for (byte[] bytes : list) {
            buffer.put(bytes);
            buffer.put(lineBytes);
        }

        byte[] data = buffer.array();
        log.info(String.format("Executing stream load to: '%s', size: '%s'", loadUrl, data.length));
        final HttpClientBuilder httpClientBuilder = HttpClients.custom()
                .setRedirectStrategy(new DefaultRedirectStrategy() {
                    @Override
                    protected boolean isRedirectable(String method) {
                        return true;
                    }
                });
        try (CloseableHttpClient httpclient = httpClientBuilder.build()) {
            HttpPut httpPut = new HttpPut(loadUrl);
            httpPut.setHeader("column_separator", ",");
            httpPut.setHeader("row_delimiter", "\r\n");
            httpPut.setHeader("columns", "tag_id,tag_name,userids,userids=to_bitmap(userids)");
            httpPut.setHeader("Expect", "100-continue");
            // httpPut.setHeader("label", label);
            httpPut.setHeader("Content-Type", "application/x-www-form-urlencoded");
            httpPut.setHeader("Authorization", getBasicAuthHeader(username, password));
            httpPut.setEntity(new ByteArrayEntity(data));
            httpPut.setConfig(RequestConfig.custom().setRedirectsEnabled(true).build());
            try (CloseableHttpResponse resp = httpclient.execute(httpPut)) {
                int code = resp.getStatusLine().getStatusCode();
                if (200 != code) {
                    log.warn("Request failed with code:{}", code);
                }
                HttpEntity respEntity = resp.getEntity();
                if (null == respEntity) {
                    log.warn("Request failed with empty response.");
                }
                log.info("result = {}", EntityUtils.toString(respEntity));
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    private String getBasicAuthHeader(String username, String password) {
        String auth = username + ":" + password;
        byte[] encodedAuth = Base64.encodeBase64(auth.getBytes(StandardCharsets.UTF_8));
        return new StringBuilder("Basic ").append(new String(encodedAuth)).toString();
    }

    public static void main(String[] args) {
        List<byte[]> list = new ArrayList<>();
        byte[] record = "4,tag2,320281199508233775".getBytes();
        byte[] record2 = "4,tag2,320281199508233776".getBytes();
        list.add(record);
        list.add(record2);

        System.out.println(6000000l / 10000);
    }
}
